566f9a
@@ -1,5 +1,5 @@
 /*
- * Copyright 2011-2012 the original author or authors.
+ * Copyright 2011-2014 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,7 +25,10 @@
import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 import org.springframework.util.Assert;
 
 import com.mongodb.DBObject;
+import com.mongodb.DBRef;
 import com.mysema.query.mongodb.MongodbSerializer;
+import com.mysema.query.types.Constant;
+import com.mysema.query.types.Operation;
 import com.mysema.query.types.Path;
 import com.mysema.query.types.PathMetadata;
 import com.mysema.query.types.PathType;
@@ -34,6 +37,7 @@
import com.mysema.query.types.PathType;
  * Custom {@link MongodbSerializer} to take mapping information into account when building keys for constraints.
  * 
  * @author Oliver Gierke
+ * @author Christoph Strobl
  */
 class SpringDataMongodbSerializer extends MongodbSerializer {
 
@@ -44,7 +48,7 @@
class SpringDataMongodbSerializer extends MongodbSerializer {
 	/**
 	 * Creates a new {@link SpringDataMongodbSerializer} for the given {@link MappingContext}.
 	 * 
-	 * @param mappingContext
+	 * @param mappingContext must not be {@literal null}.
 	 */
 	public SpringDataMongodbSerializer(MongoConverter converter) {
 
@@ -86,4 +90,57 @@
class SpringDataMongodbSerializer extends MongodbSerializer {
 
 		return super.asDBObject(key, value instanceof Pattern ? value : converter.convertToMongoType(value));
 	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see com.mysema.query.mongodb.MongodbSerializer#isReference(com.mysema.query.types.Path)
+	 */
+	@Override
+	protected boolean isReference(Path<?> path) {
+
+		MongoPersistentProperty property = getPropertyFor(path);
+		return property == null ? false : property.isAssociation();
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see com.mysema.query.mongodb.MongodbSerializer#asReference(java.lang.Object)
+	 */
+	@Override
+	protected DBRef asReference(Object constant) {
+		return converter.toDBRef(constant, null);
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see com.mysema.query.mongodb.MongodbSerializer#asReference(com.mysema.query.types.Operation, int)
+	 */
+	@Override
+	protected DBRef asReference(Operation<?> expr, int constIndex) {
+
+		for (Object arg : expr.getArgs()) {
+
+			if (arg instanceof Path) {
+
+				MongoPersistentProperty property = getPropertyFor((Path<?>) arg);
+				Object constant = ((Constant<?>) expr.getArg(constIndex)).getConstant();
+
+				return converter.toDBRef(constant, property);
+			}
+		}
+
+		return super.asReference(expr, constIndex);
+	}
+
+	private MongoPersistentProperty getPropertyFor(Path<?> path) {
+
+		Path<?> parent = path.getMetadata().getParent();
+
+		if (parent == null) {
+			return null;
+		}
+
+		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(parent.getType());
+		return entity != null ? entity.getPersistentProperty(path.getMetadata().getName()) : null;
+	}
 }
